'''
* This is the projet for Brtc LlmOps Platform
* @Author Leon-liao <liaosiliang@alltman.com>
* @Description //TODO 
* @File: 9_study_children_graph.py
* @Time: 2025/11/4
* @All Rights Reserve By Brtc
'''
import dotenv
from langchain_community.tools import GoogleSerperRun
from langchain_community.utilities import GoogleSerperAPIWrapper
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI
from langgraph.graph import MessagesState, StateGraph
from langgraph.prebuilt import ToolNode, tools_condition
from pydantic import BaseModel, Field
from typing_extensions import TypedDict, Annotated, Any
dotenv.load_dotenv()

class GoogleSerperSchema(BaseModel):
    query:str = Field(description="执行谷歌搜索的查询语句")

google_serper = GoogleSerperRun(
    name = "google_serper",
    description = ("一个低成本的谷歌搜索API工具！"),
    args_schema=GoogleSerperSchema,
    api_wrapper=GoogleSerperAPIWrapper(),)

llm = ChatOpenAI(model="gpt-4o-mini")

def reduce_str(left:str|None, right:str|None)->str:
    """归纳函数"""
    if right is not None and right != "":
        return right
    return left

class AgentState(TypedDict):
    query:Annotated[str, reduce_str]
    live_content:Annotated[str, reduce_str]
    xhs_content:Annotated[str, reduce_str]

class LiveAgentState(AgentState, MessagesState):
    """直播带货智能体状态"""
    pass

class XhsAgentState(AgentState, MessagesState):
    """小红书文案智能体文案状态"""
    pass

def chatbot_live(state:LiveAgentState)->Any:
    """直播文案智能体机器人"""
    #1、创建提示词模板 + 链应用
    prompt = ChatPromptTemplate.from_messages([
        ("system", "你是一个拥有十年经验的直播文案专家,请根据用户的提供的产品整理一篇直播带货文案，如果你的知识库找不到关于该产品信息，可以使用搜索工具"),
        ("human","{query}"),
        ("placeholder", "{chat_history}")
    ])

    chain = prompt|llm.bind_tools([google_serper])
    #2、调用链并生成AI 消息
    ai_message = chain.invoke({"query":state["query"],"chat_history":state["messages"]})
    return {
        "messages":ai_message,
        "live_content":ai_message.content
    }


# 创建直播的子图
live_agent_graph = StateGraph(LiveAgentState)
#添加节点
live_agent_graph.add_node("chatbot_live", chatbot_live)
live_agent_graph.add_node("tools", ToolNode([google_serper]))
#添加边
live_agent_graph.set_entry_point("chatbot_live")
live_agent_graph.add_conditional_edges("chatbot_live", tools_condition)
live_agent_graph.add_edge("tools", "chatbot_live")

def chatbot_xhs(state:XhsAgentState)->Any:
    """小红书文案机器人"""
    #1、创建提示词模板
    prompt = ChatPromptTemplate.from_messages([
        ("system",
         "你是一个小红书文案大师，请根据用户传递的商品名，生成一篇关于该商品的小红书笔记文案，注意风格活泼，多使用emoji表情。"),
        ("human", "{query}"),
    ])
    chain = prompt|llm|StrOutputParser()
    #2、调用链并生成内容更新状态
    return {"xhs_content":chain.invoke({"query":state["query"]})}
#创建小红的子图
xhs_agent_graph = StateGraph(XhsAgentState)
xhs_agent_graph.add_node("chatbot_xhs", chatbot_xhs)
xhs_agent_graph.set_entry_point("chatbot_xhs")
xhs_agent_graph.set_finish_point("chatbot_xhs")

# 构建大图
def parallel_node(state:AgentState)->Any:
    return state

app_graph = StateGraph(AgentState)
app_graph.add_node("parallel_node", parallel_node)
app_graph.add_node("xhs_agent", xhs_agent_graph.compile())
app_graph.add_node("live_agent",live_agent_graph.compile())

app_graph.set_entry_point("parallel_node")
app_graph.add_edge("parallel_node", "xhs_agent")
app_graph.add_edge("parallel_node", "live_agent")
app_graph.set_finish_point("xhs_agent")
app_graph.set_finish_point("live_agent")

#4、编译大图
agent = app_graph.compile()
state = agent.invoke({"query":"博睿AI培训！"})
print(state["xhs_content"])
print("===========================================")
print(state["live_content"])


